{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "import random\n",
    "import warnings\n",
    "import argparse\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import roc_curve, roc_auc_score, auc, precision_recall_curve\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from collections import defaultdict as ddict\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser= argparse.ArgumentParser(description= 'Parser for Arguments')\n",
    "parser.add_argument('-seed', type= int, default= 0)\n",
    "parser.add_argument('-num_ent', type= int, default= 1209+ 172+ 154)\n",
    "parser.add_argument('-num_drug', type= int, default= 1209)\n",
    "parser.add_argument('-num_micr', type= int, default= 172)\n",
    "parser.add_argument('-num_dise', type= int, default= 154)\n",
    "# drug_micr_rel, 0; drug_dise_rel, 1; micr_dise_rel, 2; drug_inter_rel, 3; micr_inter_rel, 4; micr_drug_rel, 5; dise_drug_rel, 6; dies_micr_rel, 7.\n",
    "parser.add_argument('-num_rel', type= int, default= 8)\n",
    "parser.add_argument('-drug_name_path', type= str, default= '../mdd/drug/drug_name.txt')\n",
    "parser.add_argument('-micr_name_path', type= str, default= '../mdd/microbe/microbe_name.txt')\n",
    "parser.add_argument('-dise_name_path', type= str, default= '../mdd/disease/disease_name.txt')\n",
    "parser.add_argument('-drug_micr_adj_path', type= str, default= '../mdd/adj/microbe_drug_adj.txt')\n",
    "parser.add_argument('-drug_struct_simi_path', type= str, default= '../mdd/drug/drug_struct_simi.txt')\n",
    "parser.add_argument('-drug_inter_adj_path', type= str, default= '../mdd/drug/drug_interact_adj.txt')\n",
    "parser.add_argument('-drug_dise_adj_path', type= str, default= '../mdd/adj/drug_disease_adj.txt')\n",
    "parser.add_argument('-micr_ani_path', type= str, default= '../mdd/microbe/microbe_ani_simi.txt')\n",
    "parser.add_argument('-micr_inter_adj_path', type= str, default= '../mdd/microbe/microbe_interact_adj.txt')\n",
    "parser.add_argument('-micr_gene_simi_path', type= str, default= '../mdd/microbe/microbe_gene_simi.txt')\n",
    "parser.add_argument('-micr_dise_adj_path', type= str, default= '../mdd/adj/microbe_disease_adj.txt')\n",
    "parser.add_argument('-dise_simi_path', type= str, default= '../mdd/disease/disease_dag_simi.txt')\n",
    "parser.add_argument('-train_ratio', type= float, default= 0.8)\n",
    "parser.add_argument('-valid_ratio', type= float, default= 0.1)\n",
    "parser.add_argument('-test_ratio', type= float, default= 0.1)\n",
    "parser.add_argument('-kg_file', type= str, default= 'kg_data/')\n",
    "parser.add_argument('-batch_size', type= int, default= 128)\n",
    "parser.add_argument('-lbl_smooth', type= float, default= 0.2, help= 'Label smoothing enable or disable')\n",
    "parser.add_argument('-embed_dim', type= int, default= 128)\n",
    "parser.add_argument('-device', type= str, default= 'cuda:0')\n",
    "parser.add_argument('-lr_net1', type= float, default= 1e-3)\n",
    "parser.add_argument('-weight_decay_net1', type= float, default= 0)\n",
    "parser.add_argument('-patience_net1', type= int, default= 50)\n",
    "parser.add_argument('-epoch_net1', type= int, default= 10000)\n",
    "parser.add_argument('-lr_net2', type= float, default= 1e-3)\n",
    "parser.add_argument('-weight_decay_net2', type= float, default= 1e-4)\n",
    "parser.add_argument('-patience_net2', type= int, default= 6)\n",
    "parser.add_argument('-epoch_net2', type= int, default= 300)\n",
    "parser.add_argument('-pt_file', type= str, default= 'checkpoint/')\n",
    "parser.add_argument('-balance_factor_net1', type= float, default= 1)\n",
    "parser.add_argument('-threshold_net1', type= float, default= 0.8)\n",
    "parser.add_argument('-walk_prob', type= float, default= 0.8)\n",
    "parser.add_argument('-walk_epoch', type= int, default= 500)\n",
    "parser.add_argument('-walk_eps', type= float, default= 1e-10)\n",
    "parser.add_argument('-walk_err', type= float, default= 1e-6)\n",
    "parser.add_argument('-pt_file_name_net1', type= str, default= 'net1s1.pt')\n",
    "parser.add_argument('-pt_file_name_net2', type= str, default= 'net2s1.pt')\n",
    "parser.add_argument('-test_result_file', type= str, default= '../result/pcmda_test_result.txt')\n",
    "params= parser.parse_args([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class dataloader(object):\n",
    "    def __init__(self, params):\n",
    "        super().__init__()\n",
    "        self.params= params\n",
    "        self.drug_micr_asso_mat, self.drug_dise_asso_mat, self.drug_inter_mat, self.drug_struct_simi_mat= self.load_drug_data()\n",
    "        self.train_xy, self.valid_xy, self.test_xy= self.split_dataset()\n",
    "        self.drug_micr_asso_mat_zy= self.get_asso_mat_zy()\n",
    "        self.micr_ani_mat, self.micr_inter_mat, self.micr_dise_asso_mat, self.micr_asso_simi_mat, self.micr_gene_simi_mat, self.micr_inte_simi_mat, self.drug_asso_simi_mat= self.load_micr_data()\n",
    "        self.dise_simi_mat, self.drug_dise_drug_simi_mat, self.micr_dise_micr_simi_mat= self.load_dise_data()\n",
    "        # to ensure less noise, we use consistence matrix instead of microbe integate similarity matrix to construct hete graph.\n",
    "        self.hete_graph_mat= torch.cat([torch.cat([self.drug_struct_simi_mat, self.drug_micr_asso_mat_zy, self.drug_dise_asso_mat], dim= 1),\\\n",
    "                                    torch.cat([self.drug_micr_asso_mat_zy.T, self.micr_ani_mat, self.micr_dise_asso_mat], dim= 1),\\\n",
    "                                    torch.cat([self.drug_dise_asso_mat.T, self.micr_dise_asso_mat.T, self.dise_simi_mat], dim= 1)], dim= 0).float()\n",
    "        self.struct_coef_mat1= self.random_walk_root(self.hete_graph_mat, params.walk_prob, params.walk_epoch, 1, params.walk_eps, params.walk_err)\n",
    "        self.struct_coef_mat2= self.random_walk_root(self.hete_graph_mat, params.walk_prob, params.walk_epoch, 2, params.walk_eps, params.walk_err)\n",
    "        self.struct_coef_mat3= self.random_walk_root(self.hete_graph_mat, params.walk_prob, params.walk_epoch, 3, params.walk_eps, params.walk_err)\n",
    "        self.create_kg_data()\n",
    "        self.introduce()\n",
    "\n",
    "    # @introduce data\n",
    "    def introduce(self):\n",
    "        print(f'Drug microbe association num: {self.drug_micr_asso_mat.sum()}\\nDrug interaction num: {self.drug_inter_mat.sum()}\\nMicrobe interaction num: {self.micr_inter_mat.sum()}')\n",
    "        print(f'Drug disease association num: {self.drug_dise_asso_mat.sum()}\\nMicrobe disease association num: {self.micr_dise_asso_mat.sum()}')\n",
    "\n",
    "    # @ mask\n",
    "    def get_asso_mat_zy(self):\n",
    "        asso_mat_zy= self.drug_micr_asso_mat.clone()\n",
    "        asso_mat_zy[self.valid_xy[:, 0], self.valid_xy[:, 1]]= 0\n",
    "        asso_mat_zy[self.test_xy[:, 0], self.test_xy[:, 1]]= 0\n",
    "        return asso_mat_zy\n",
    "\n",
    "    # @ 计算hamming interaction profile similarity.\n",
    "    def hip_sim(self, mat):\n",
    "        sim_ls, dim= [], mat.shape[1]\n",
    "        for i in range(mat.shape[0]):\n",
    "            sim_ls.append(((mat[i]- mat) == 0).sum(dim= 1)/ dim)\n",
    "        return torch.stack(sim_ls)\n",
    "    \n",
    "    # norm1\n",
    "    def norm_z(self, x, eps= 1e-10):\n",
    "        return (x- x.mean())/ (x.std()+ eps)\n",
    "\n",
    "    # norm2\n",
    "    def norm_min_max(self, x, eps= 1e-10):\n",
    "        return (x- x.min())/ (x.max()- x.min()+ eps)    \n",
    "    \n",
    "    # @create knowledge graph data\n",
    "    def create_kg_data(self):\n",
    "        if os.path.exists(self.params.kg_file)== False:os.makedirs(self.params.kg_file)\n",
    "        # drug microbe association data\n",
    "        drug_micr_train, drug_micr_valid, drug_micr_test= self.train_xy[self.drug_micr_asso_mat[self.train_xy[:, 0], self.train_xy[:, 1]]== 1],\\\n",
    "        self.valid_xy[self.drug_micr_asso_mat[self.valid_xy[:, 0], self.valid_xy[:, 1]]== 1],\\\n",
    "        self.test_xy[self.drug_micr_asso_mat[self.test_xy[:, 0], self.test_xy[:, 1]]== 1]\n",
    "        # drug disease association data\n",
    "        drug_dise_train, micr_dise_train, drug_inter_train, micr_inter_train= self.drug_dise_asso_mat.nonzero(), self.micr_dise_asso_mat.nonzero(), self.drug_inter_mat.nonzero(), self.micr_inter_mat.nonzero()\n",
    "        # add offset\n",
    "        drug_micr_valid+= torch.tensor([0, self.params.num_drug]); drug_micr_test+= torch.tensor([0, self.params.num_drug]); drug_micr_train+= torch.tensor([0, self.params.num_drug])\n",
    "        drug_dise_train+= torch.tensor([0, self.params.num_drug+ self.params.num_micr]); micr_dise_train+= torch.tensor([self.params.num_drug, self.params.num_drug+ self.params.num_micr]); micr_inter_train+= torch.tensor([self.params.num_drug, self.params.num_drug])\n",
    "        # add rel\n",
    "        drug_micr_train, drug_micr_valid, drug_micr_test= drug_micr_train[:, [0, 1, 1]], drug_micr_valid[:, [0, 1, 1]], drug_micr_test[:, [0, 1, 1]]\n",
    "        drug_dise_train, micr_dise_train, drug_inter_train, micr_inter_train= drug_dise_train[:, [0, 1, 1]], micr_dise_train[:, [0, 1, 1]], drug_inter_train[:, [0, 1, 1]], micr_inter_train[:, [0, 1, 1]]\n",
    "        drug_micr_train[:, 2], drug_micr_valid[:, 2], drug_micr_test[:, 2], drug_dise_train[:, 2], micr_dise_train[:, 2], drug_inter_train[:, 2], micr_inter_train[:, 2]= 0, 0, 0, 1, 2, 3, 4\n",
    "        # savefile\n",
    "        train= torch.cat([drug_micr_train, drug_dise_train, micr_dise_train, drug_inter_train, micr_inter_train], dim= 0)\n",
    "        val= drug_micr_valid\n",
    "        test= drug_micr_test\n",
    "        np.savetxt(f'{params.kg_file}//train.txt', train, fmt= '%d', delimiter= '\\t', encoding= 'utf-8-sig')\n",
    "        np.savetxt(f'{params.kg_file}//valid.txt', val, fmt= '%d', delimiter= '\\t', encoding= 'utf-8-sig')\n",
    "        np.savetxt(f'{params.kg_file}//test.txt', test, fmt= '%d', delimiter= '\\t', encoding= 'utf-8-sig')\n",
    "        print(f'Knowledge graph data has prepared...')\n",
    "\n",
    "    # @split data set\n",
    "    def split_dataset(self):\n",
    "        train_xy, valid_xy, test_xy= [], [], []\n",
    "        for i in range(self.params.num_drug):\n",
    "            first= True\n",
    "            for j in range(self.params.num_micr):\n",
    "                if self.drug_micr_asso_mat[i, j]== 1 and first:\n",
    "                    train_xy.append([i, j])\n",
    "                    first= False\n",
    "                else:\n",
    "                    num= torch.rand(1)\n",
    "                    if num< self.params.train_ratio:\n",
    "                        train_xy.append([i, j])\n",
    "                    elif num>= self.params.train_ratio and num< self.params.train_ratio+ self.params.valid_ratio:\n",
    "                        valid_xy.append([i, j])\n",
    "                    else:\n",
    "                        test_xy.append([i, j])        \n",
    "        print(f'Spliting data has finished...')\n",
    "        return torch.tensor(train_xy), torch.tensor(valid_xy), torch.tensor(test_xy)\n",
    "\n",
    "    # @load disease data\n",
    "    def load_dise_data(self):\n",
    "        dise_simi_mat= torch.from_numpy(np.loadtxt(self.params.dise_simi_path, encoding= 'utf-8-sig'))\n",
    "        drug_dise_drug_simi_mat, micr_dise_micr_simi_mat= torch.matmul(nn.functional.normalize(self.drug_inter_mat, p= 2, dim= 1), nn.functional.normalize(self.drug_inter_mat, p= 2, dim= 1).T),\\\n",
    "        torch.matmul(nn.functional.normalize(self.micr_inter_mat, p= 2, dim= 1), nn.functional.normalize(self.micr_inter_mat, p= 2, dim= 1).T)\n",
    "        for i in range(self.params.num_drug): drug_dise_drug_simi_mat[i, i]= 1.0\n",
    "        for i in range(self.params.num_micr): micr_dise_micr_simi_mat[i, i]= 1.0\n",
    "        return dise_simi_mat, drug_dise_drug_simi_mat, micr_dise_micr_simi_mat\n",
    "\n",
    "    # @load micr data\n",
    "    def load_micr_data(self):\n",
    "        micr_ani_mat= torch.from_numpy(np.loadtxt(self.params.micr_ani_path, encoding= 'utf-8-sig'))/ 100\n",
    "        micr_gene_simi_mat= (torch.from_numpy(np.loadtxt(self.params.micr_gene_simi_path, encoding= 'utf-8-sig'))+ 1)/ 2\n",
    "        micr_inter_mat= self.load_adj_data(self.params.micr_inter_adj_path, sp= (self.params.num_micr, self.params.num_micr))\n",
    "        micr_dise_mat= self.load_adj_data(self.params.micr_dise_adj_path, sp= (self.params.num_micr, self.params.num_dise))\n",
    "        micr_asso_simi_mat= torch.matmul(nn.functional.normalize(self.drug_micr_asso_mat_zy.T, p= 2, dim= 1), nn.functional.normalize(self.drug_micr_asso_mat_zy.T, p= 2, dim= 1).T)\n",
    "        drug_asso_simi_mat= torch.matmul(nn.functional.normalize(self.drug_micr_asso_mat_zy, p= 2, dim= 1), nn.functional.normalize(self.drug_micr_asso_mat_zy, p= 2, dim= 1).T)\n",
    "        for i in range(self.params.num_drug):drug_asso_simi_mat[i, i]= 1\n",
    "        for i in range(self.params.num_micr):micr_asso_simi_mat[i, i]= 1\n",
    "        for i in range(self.params.num_micr):micr_ani_mat[i, i]= 1\n",
    "        return micr_ani_mat, micr_inter_mat, micr_dise_mat, micr_asso_simi_mat, micr_gene_simi_mat, torch.where(micr_ani_mat> 0, micr_ani_mat, micr_gene_simi_mat), drug_asso_simi_mat\n",
    "\n",
    "    # @load drug data\n",
    "    def load_drug_data(self):\n",
    "        drug_dise_asso_mat= self.load_adj_data(self.params.drug_dise_adj_path, sp= (self.params.num_drug, self.params.num_dise))\n",
    "        drug_micr_asso_mat= self.load_adj_data(self.params.drug_micr_adj_path, sp= (self.params.num_drug, self.params.num_micr))\n",
    "        drug_inter_mat= self.load_adj_data(self.params.drug_inter_adj_path, sp= (self.params.num_drug, self.params.num_drug))\n",
    "        drug_struct_simi_mat= torch.from_numpy(np.loadtxt(self.params.drug_struct_simi_path, encoding= 'utf-8-sig'))\n",
    "        return drug_micr_asso_mat, drug_dise_asso_mat, drug_inter_mat, drug_struct_simi_mat\n",
    "    \n",
    "    # @load adj data\n",
    "    def load_adj_data(self, path, sp= (1209, 172)):\n",
    "        idx= torch.from_numpy(np.loadtxt(path, encoding= 'utf-8-sig')).long()- 1\n",
    "        mat= torch.zeros((sp[0], sp[1]))\n",
    "        mat[idx[:, 0], idx[:, 1]]= 1\n",
    "        return mat\n",
    "    \n",
    "    # @NRWR, neighborhood random walk with restart\n",
    "    def random_walk_root(self, A, alpha, epoch, nei, eps, err):\n",
    "        Mask_mat= torch.sign(torch.matrix_power(A, nei)- torch.matrix_power(A, nei- 1))\n",
    "        W= nn.functional.normalize(A, p= 1, dim= 1)\n",
    "        S_last= torch.eye(A.shape[0])\n",
    "        for i in range(epoch):\n",
    "            S_new= alpha* torch.matmul(S_last, W)+ (1- alpha)* torch.eye(A.shape[0])\n",
    "            S_new= nn.functional.normalize(Mask_mat* S_new, p= 1, dim= 1)\n",
    "            if (S_new- S_last).abs().sum()<= err* A.shape[0]** 2:\n",
    "                print('converge...');break\n",
    "            S_last= S_new\n",
    "        return S_new\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EarlyStopping:\n",
    "\t\"\"\"docstring for EarlyStopping\"\"\"\n",
    "\tdef __init__(self, patience, pt_file= 'checkpoint/', file_name= 'checkpoint.pt', mess_out= True, eps= 0):\n",
    "\t\tsuper().__init__()\n",
    "\t\tself.patience, self.eps, self.pt_file, self.file_name, self.mess_out= patience, eps, pt_file, file_name, mess_out\n",
    "\t\tself.best_score, self.counter, self.flag= None, 0, False\n",
    "\t\tif os.path.exists(self.pt_file)== False:os.makedirs(self.pt_file)\n",
    "\t\n",
    "\tdef __call__(self, val_loss, model):\n",
    "\t\tscore= -val_loss\n",
    "\t\tif self.best_score is None:\n",
    "\t\t\tself.best_score= score\n",
    "\t\t\tself.save_checkpoint(model)\n",
    "\t\telif score<= self.best_score- self.eps:\n",
    "\t\t\tself.counter+= 1\n",
    "\t\t\tif self.mess_out:print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
    "\t\t\tif self.counter>= self.patience:\n",
    "\t\t\t\tself.flag= True\n",
    "\t\telse:\n",
    "\t\t\tself.best_score= score\n",
    "\t\t\tself.save_checkpoint(model)\n",
    "\t\t\tself.counter= 0\n",
    "\n",
    "\tdef save_checkpoint(self, model):\n",
    "\t\ttorch.save(model, f'{self.pt_file}//{self.file_name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MDDataset(Dataset):    \n",
    "    def __init__(self, triples, split, params):\n",
    "        self.triples= triples\n",
    "        self.split= split\n",
    "        self.params= params\n",
    "    def __len__(self):\n",
    "        return len(self.triples)\n",
    "    def __getitem__(self, idx):\n",
    "        ele= self.triples[idx]\n",
    "        triple, label= torch.LongTensor(ele['triple']), np.int32(ele['label'])\n",
    "        label= self.get_label(label)\n",
    "        if self.split== 'train' and self.params.lbl_smooth!= 0.0:\n",
    "            label= (1.0- self.params.lbl_smooth)* label+ (1.0/ self.params.num_ent)\n",
    "        return triple, label\n",
    "    @staticmethod\n",
    "    def collate_fn(data):\n",
    "        triple= torch.stack([_[0] for _ in data], dim= 0)\n",
    "        label= torch.stack([_[1] for _ in data], dim= 0)\n",
    "        return triple, label\n",
    "    def get_label(self, label):\n",
    "        y= np.zeros([self.params.num_ent], dtype= np.float32)\n",
    "        for e2 in label: y[e2]= 1.0\n",
    "        return torch.FloatTensor(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(params):\n",
    "    sr2d, data, tp= ddict(set), ddict(list), 0\n",
    "    for split in ['train', 'valid', 'test']:\n",
    "        for line in open(f'{params.kg_file}{split}.txt', encoding= 'utf-8-sig'):\n",
    "            src_id, dst_id, rel_id= line.strip().split('\\t')\n",
    "            src_id, rel_id, dst_id= int(float(src_id)), int(float(rel_id)), int(float(dst_id))\n",
    "            data[split].append((src_id, rel_id, dst_id))\n",
    "            if split== 'train':\n",
    "                sr2d[(src_id, rel_id)].add(dst_id)\n",
    "                if rel_id in [0, 1, 2]:sr2d[(dst_id, rel_id+ 5)].add(src_id)\n",
    "    sr2d4tr= {k: list(v) for k, v in sr2d.items()}\n",
    "    triples= ddict(list)\n",
    "    for (src_id, rel_id), dst_id in sr2d4tr.items():\n",
    "        triples['train'].append({'triple': (src_id, rel_id, -1), 'label': dst_id})\n",
    "    for split in ['test', 'valid']:\n",
    "        for src_id, rel_id, dst_id in data[split]:\n",
    "            sr2d[(src_id, rel_id)].add(dst_id)\n",
    "            sr2d[(dst_id, rel_id+ 5)].add(src_id)\n",
    "    sr2d4val_te= {k: list(v) for k, v in sr2d.items()}\n",
    "    for split in ['valid', 'test']:\n",
    "        for src_id, rel_id, dst_id in data[split]:\n",
    "            triples[f'{split}'].append({'triple': (src_id, rel_id, dst_id), 'label': sr2d4val_te[(src_id, rel_id)]})\n",
    "    triples= dict(triples)\n",
    "    def get_data_loader(dataset_class, split, batch_size, shuffle= True):\n",
    "        return DataLoader(dataset_class(triples[split], split, params), batch_size= batch_size, shuffle= shuffle, collate_fn= dataset_class.collate_fn)\n",
    "    data_iter= {\n",
    "        'train': get_data_loader(MDDataset, 'train', params.batch_size),\n",
    "        'valid': get_data_loader(MDDataset, 'valid', params.batch_size),\n",
    "        'test': get_data_loader(MDDataset, 'test', params.batch_size)}\n",
    "    return data_iter, triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAKG(torch.nn.Module):\n",
    "    def __init__(self, params, struct_coef_mat1, struct_coef_mat2):\n",
    "        super().__init__()\n",
    "        self.params= params\n",
    "        self.scm1, self.scm2= nn.functional.normalize(struct_coef_mat1, p= 2, dim= 1), nn.functional.normalize(struct_coef_mat2, p= 2, dim= 1)\n",
    "        self.fc4att= nn.Sequential(nn.Linear(self.params.embed_dim* 3, self.params.embed_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(self.params.embed_dim, 2))\n",
    "        self.fc1, self.fc1res= nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Dropout(0.3), nn.Linear(16, 1)), nn.Linear(2, 1)\n",
    "        self.fc2= nn.Sequential(nn.Linear(self.params.embed_dim, 2* self.params.embed_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(2* self.params.embed_dim, self.params.embed_dim))\n",
    "        self.ent_embed= torch.nn.Embedding(self.params.num_ent, self.params.embed_dim, padding_idx= None); nn.init.xavier_normal_(self.ent_embed.weight)\n",
    "        self.rel_embed= torch.nn.Embedding(self.params.num_rel, self.params.embed_dim, padding_idx= None); nn.init.xavier_normal_(self.rel_embed.weight)\n",
    "        self.bceloss= torch.nn.BCELoss()\n",
    "        self.init_para()\n",
    "    def init_para(self):\n",
    "        nn.init.xavier_normal_(self.fc4att[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc4att[3].weight)\n",
    "        nn.init.xavier_normal_(self.fc1[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc1[3].weight)\n",
    "        nn.init.xavier_normal_(self.fc1res.weight)\n",
    "        nn.init.xavier_normal_(self.fc2[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc2[3].weight)\n",
    "    def loss(self, pred, true_label, mats, balance_factor= 0.3, threshold= 0.8):\n",
    "        drug_embed, micr_embed, dise_embed= nn.functional.normalize(self.ent_embed.weight[0: self.params.num_drug], p= 2, dim= 1), nn.functional.normalize(self.ent_embed.weight[self.params.num_drug: self.params.num_drug+ self.params.num_micr], p= 2, dim= 1), nn.functional.normalize(self.ent_embed.weight[self.params.num_drug+ self.params.num_micr: ], p= 2, dim= 1)\n",
    "        mat0, mat1, mat2= torch.matmul(drug_embed, drug_embed.T), torch.matmul(micr_embed, micr_embed.T), torch.matmul(dise_embed, dise_embed.T)\n",
    "        mse_loss= (((mat0- mats[0])* ((mats[0]!= 0)| (mat0>= threshold)))** 2).mean()+ ((mat1- mats[1])** 2).mean()+ ((mat2- mats[2])** 2).mean()\n",
    "        return balance_factor* mse_loss+ self.bceloss(pred, true_label)\n",
    "    def forward(self, src, rel):\n",
    "        X1, X2, X3= self.ent_embed.weight, self.scm1@ self.ent_embed.weight, self.scm2@ self.ent_embed.weight\n",
    "        att_mat= torch.softmax(self.fc4att(torch.cat([X1, X2, X3], dim= 1)), dim= 1)\n",
    "        X= X1+ att_mat[:, 0].view(-1, 1)* X2+ att_mat[:, 1].view(-1, 1)* X3\n",
    "        src_emb, rel_emb= X[src], self.rel_embed(rel)\n",
    "        fea1= torch.cat([src_emb.unsqueeze(1), rel_emb.unsqueeze(1)], dim= 1).transpose(1, 2)\n",
    "        fea2= (self.fc1(fea1)+ self.fc1res(fea1)).squeeze(2)\n",
    "        return torch.sigmoid(torch.matmul(self.fc2(fea2), X.T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(net, data_iter, params, split= 'valid'):\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        results= {}\n",
    "        train_iter= iter(data_iter[split])\n",
    "        for step, batch in enumerate(train_iter):\n",
    "            triple, label= [_.to(params.device) for _ in batch]\n",
    "            src, rel, dst, label= triple[:, 0], triple[:, 1], triple[:, 2], label\n",
    "            pred= net.forward(src, rel)\n",
    "            b_range= torch.arange(pred.size()[0], device= params.device)\n",
    "            target_pred= pred[b_range, dst]\n",
    "            pred= torch.where(label.byte(), torch.zeros_like(pred), pred)\n",
    "            pred[b_range, dst]= target_pred\n",
    "            pred= pred.cpu().numpy()\n",
    "            dst= dst.cpu().numpy()\n",
    "            for i in range(pred.shape[0]):\n",
    "                scores= pred[i]\n",
    "                target= dst[i]\n",
    "                tar_scr= scores[target]\n",
    "                scores= np.delete(scores, target)\n",
    "                rand= np.random.randint(scores.shape[0])\n",
    "                scores= np.insert(scores, rand, tar_scr)\n",
    "                sorted_indices= np.argsort(-scores, kind= 'stable')\n",
    "                _filter= np.where(sorted_indices== rand)[0][0]\n",
    "                results['count']= 1+ results.get('count', 0.0)\n",
    "                results['mr']= (_filter+ 1)+ results.get('mr', 0.0)\n",
    "                results['mrr']= (1.0/ (_filter+ 1))+ results.get('mrr', 0.0)\n",
    "                for k in range(10):\n",
    "                    if _filter<= k:\n",
    "                        results[f'hits@{k+ 1}']= 1+ results.get(f'hits@{k+ 1}', 0.0)\n",
    "    results['mr']= round(results['mr']/ float(results['count']), 5)\n",
    "    results['mrr']= round(results['mrr']/ float(results['count']), 5)\n",
    "    for k in range(10):\n",
    "        results[f'hits@{k+1}']= round(results.get(f'hits@{k+ 1}', 0)/ float(results['count']), 5)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(params.seed)\n",
    "np.random.seed(params.seed)\n",
    "torch.manual_seed(params.seed)\n",
    "torch.cuda.manual_seed(params.seed)\n",
    "torch.cuda.manual_seed_all(params.seed)\n",
    "os.environ['PYTHONHASHSEED'] = str(params.seed)    \n",
    "torch.backends.cudnn.deterministic = True\n",
    "dl= dataloader(params)\n",
    "# split dataset\n",
    "train_label, valid_label, test_label= dl.drug_micr_asso_mat[dl.train_xy[:, 0], dl.train_xy[:, 1]].long(), dl.drug_micr_asso_mat[dl.valid_xy[:, 0], dl.valid_xy[:, 1]].long(), dl.drug_micr_asso_mat[dl.test_xy[:, 0], dl.test_xy[:, 1]].long()\n",
    "train_xy_label_dataset= torch.utils.data.TensorDataset(dl.train_xy, train_label)\n",
    "train_loader= torch.utils.data.DataLoader(train_xy_label_dataset, batch_size= params.batch_size, shuffle= True)\n",
    "valid_xy_label_dataset= torch.utils.data.TensorDataset(dl.valid_xy, valid_label)\n",
    "valid_loader= torch.utils.data.DataLoader(valid_xy_label_dataset, batch_size= params.batch_size, shuffle= False)\n",
    "test_xy_label_dataset= torch.utils.data.TensorDataset(dl.test_xy, test_label)\n",
    "test_loader= torch.utils.data.DataLoader(test_xy_label_dataset, batch_size= params.batch_size, shuffle= False)\n",
    "# knowledge graph data\n",
    "data_iter, triples= load_data(params)\n",
    "net= SAKG(params, dl.struct_coef_mat1.to(params.device), dl.struct_coef_mat2.to(params.device)).to(params.device)\n",
    "optimizer= torch.optim.Adam(net.parameters(), lr= params.lr_net1, weight_decay= params.weight_decay_net1)\n",
    "earlystopping4kg= EarlyStopping(patience= params.patience_net1, pt_file= params.pt_file, file_name= params.pt_file_name_net1, mess_out= True)\n",
    "mats= [dl.drug_struct_simi_mat.clone().to(params.device), (dl.micr_inte_simi_mat).clone().to(params.device), dl.dise_simi_mat.clone().to(params.device)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(params.epoch_net1):\n",
    "    net.train()\n",
    "    losses= []\n",
    "    train_iter= iter(data_iter['train'])\n",
    "    for step, batch in enumerate(train_iter):\n",
    "        optimizer.zero_grad()\n",
    "        triple, label= [_.to(params.device) for _ in batch]\n",
    "        src, rel, dst, label= triple[:, 0], triple[:, 1], triple[:, 2], label\n",
    "        pred= net.forward(src, rel)\n",
    "        loss= net.loss(pred, label, mats, params.balance_factor_net1, params.threshold_net1)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        losses.append(loss.item())\n",
    "        if step% 5000== 0:\n",
    "            print(f'epoch: {epoch}, Train loss: {np.mean(losses)}, rel ratio: {((rel== 0)).sum()/ len(rel)}, acc: {(torch.where(pred< 0.5, 0, 1)== torch.where(label< 0.5, 0, 1)).sum()/ label.shape[0]/ label.shape[1]}')\n",
    "    results= evaluate(net, data_iter, params, 'valid')\n",
    "    print(f'val mr: {results[\"mr\"]}, mrr: {results[\"mrr\"]}, hits10: {results[\"hits@10\"]}')\n",
    "    earlystopping4kg(-(results['mrr']+ results['hits@10']), net)\n",
    "    if earlystopping4kg.flag== True: break\n",
    "# kg valid\n",
    "net= torch.load(f'{params.pt_file}//{params.pt_file_name_net1}')\n",
    "print(f\"valid results, {evaluate(net, data_iter, params, 'valid')}\\ntest results, {evaluate(net, data_iter, params, 'test')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Conv2d(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch, ksa, ksb):\n",
    "        super().__init__()\n",
    "        assert ksb% 2== 1, 'kernel_size shoud be odd.'\n",
    "        self.ksb= ksb\n",
    "        self.conv= nn.Conv2d(in_ch, out_ch, kernel_size= (ksa, ksb))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_pad= F.pad(x, [(self.ksb- 1)// 2, (self.ksb- 1)// 2, 0, 0])\n",
    "        return self.conv(x_pad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CrossEmbedLayer(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch, ksa= 2, ksbs= [1, 3, 7, 11], stride= 1):\n",
    "        super().__init__()\n",
    "        assert stride== 1, 'It is under development.'\n",
    "        kernel_sizes= sorted(ksbs)\n",
    "        num_scales= len(kernel_sizes)\n",
    "        # calculate the dimension at each scale\n",
    "        channel_scales= [int(out_ch / (2 ** i)) for i in range(1, num_scales)]\n",
    "        channel_scales= [*channel_scales, out_ch- sum(channel_scales)]\n",
    "        # conv\n",
    "        self.convs= nn.ModuleList([])\n",
    "        for ks, ch_scale in zip(ksbs, channel_scales):\n",
    "            self.convs.append(Conv2d(in_ch, ch_scale, ksa= ksa, ksb= ks))\n",
    "        self.conv_end= nn.Conv2d(out_ch, out_ch, 1)\n",
    "    def forward(self, x):\n",
    "        fmaps= tuple(map(lambda conv: conv(x), self.convs))\n",
    "        return self.conv_end(torch.cat(fmaps, dim = 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MPFI(torch.nn.Module):\n",
    "    def __init__(self, params, ent_emb, mats):\n",
    "        super().__init__()\n",
    "        self.params, self.asso_mat, self.drug_inte_mat, self.drug_simi_mat= params, mats[0], mats[1], mats[2]\n",
    "        self.drug_emb, self.micr_emb= ent_emb[0: self.params.num_drug], ent_emb[self.params.num_drug: self.params.num_drug+ self.params.num_micr]\n",
    "        self.att1drug= nn.Parameter(torch.ones((params.embed_dim, params.embed_dim)))\n",
    "        self.fc4fuse= nn.Linear(5, 1)\n",
    "    def init_para(self):\n",
    "        nn.init.xavier_normal_(self.att1drug);nn.init.xavier_normal_(self.fc4fuse.weight)\n",
    "    def forward(self, x, y):\n",
    "        att_mat4drug= self.fc4fuse(torch.cat([(self.drug_emb[x]@ self.att1drug@ self.drug_emb.T).unsqueeze(-1), self.drug_simi_mat[x].unsqueeze(-1), self.drug_inte_mat[x].unsqueeze(-1), (self.drug_inte_mat[x]@ self.drug_inte_mat.T).unsqueeze(-1)/ self.params.num_drug, (self.drug_simi_mat[x]@ self.drug_simi_mat.T).unsqueeze(-1)/ self.params.num_drug], dim= -1)).squeeze()\n",
    "        att_mat4drug[(self.asso_mat.T)[y]== 0]= -1e10\n",
    "        att_mat4drug_= torch.cat([torch.ones((x.shape[0], 1)).to(self.params.device), att_mat4drug], dim= 1)\n",
    "        att_mat4drug_= torch.softmax(att_mat4drug_, dim= 1)\n",
    "        return torch.cat([self.drug_emb[x], self.micr_emb[y], self.drug_emb[x]- att_mat4drug_[:, 1:]@ self.drug_emb], dim= 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BCE(torch.nn.Module):\n",
    "    def __init__(self, params):\n",
    "        super().__init__()\n",
    "        self.params= params\n",
    "        self.cv4fea1= nn.Sequential(CrossEmbedLayer(3, 256, ksa= 2, stride= 1), nn.ReLU(), nn.Dropout(0.5))\n",
    "        self.fc4fea1= nn.Sequential(nn.Linear(1381, 1381), nn.ReLU(), nn.Dropout(0.5))\n",
    "        self.fc4fea3= nn.Sequential(nn.Linear(1381, 256), nn.ReLU(), nn.Dropout(0.3))\n",
    "        self.fc4gt= nn.Sequential(nn.Linear(512, 256), nn.Sigmoid(), nn.Dropout(0.3))\n",
    "        self.init_para()\n",
    "    def init_para(self):\n",
    "        # nn.init.xavier_normal_(self.cv4fea1[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc4fea1[0].weight, nn.init.calculate_gain('relu'))  \n",
    "        nn.init.xavier_normal_(self.fc4fea3[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc4gt[0].weight, nn.init.calculate_gain('sigmoid'))\n",
    "    def forward(self, fea):\n",
    "        fea= self.cv4fea1(fea)\n",
    "        fea1= nn.functional.normalize(self.fc4fea1(fea.sum(dim=  1, keepdims= True)), p= 2, dim= -1)* fea+ fea\n",
    "        fea2= nn.functional.normalize(fea1.sum(-1).squeeze(), p= 2, dim= -1)\n",
    "        fea3= nn.functional.normalize(self.fc4fea3(fea1.sum(1).squeeze()), p= 2, dim= -1)\n",
    "        gt= self.fc4gt(torch.cat([fea2, fea3], dim= 1))\n",
    "        fea4= gt* fea2+ (1- gt)* fea3\n",
    "        return fea4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PCMDA(torch.nn.Module):\n",
    "    def __init__(self, params, ent_emb, mats):\n",
    "        super().__init__()\n",
    "        self.params= params\n",
    "        self.bce, self.mpfi= BCE(params), MPFI(params, ent_emb, mats)\n",
    "        self.fc4gt= nn.Sequential(nn.Linear(5* self.params.embed_dim, self.params.embed_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(self.params.embed_dim, 2))\n",
    "        self.loss_fuc= nn.CrossEntropyLoss(reduction= 'mean')\n",
    "    def init_para(self):\n",
    "        nn.init.xavier_normal_(self.fc4gt[0].weight, nn.init.calculate_gain('relu'))\n",
    "        nn.init.xavier_normal_(self.fc4gt[3].weight)        \n",
    "    def loss(self, pred, label):\n",
    "        return self.loss_fuc(pred, label)\n",
    "    def forward(self, x, y, fea):\n",
    "        o1= self.bce(fea)\n",
    "        o2= self.mpfi(x, y)\n",
    "        return self.fc4gt(torch.cat([o1, o2], dim= -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare data\n",
    "net= torch.load(f'{params.pt_file}//{params.pt_file_name_net1}')\n",
    "X1, X2, X3= net.ent_embed.weight, net.scm1@ net.ent_embed.weight, net.scm2@ net.ent_embed.weight\n",
    "att_mat= torch.softmax(net.fc4att(torch.cat([X1, X2, X3], dim= 1)), dim= 1)\n",
    "ent_emb= (X1+ att_mat[:, 0].view(-1, 1)* X2+ att_mat[:, 1].view(-1, 1)* X3).detach().clone()\n",
    "drug_inte_mat, drug_struct_simi_mat=  dl.drug_inter_mat, dl.drug_struct_simi_mat\n",
    "mats= [dl.drug_micr_asso_mat_zy, drug_inte_mat, drug_struct_simi_mat]\n",
    "mats= [mat.to(torch.float).to(params.device) for mat in mats]\n",
    "# define network\n",
    "net= PCMDA(params, ent_emb.to(params.device), mats).to(params.device)\n",
    "optimizer= torch.optim.Adam(net.parameters(), lr= params.lr_net2, weight_decay= params.weight_decay_net2)\n",
    "earlystopping= EarlyStopping(patience= params.patience_net2, pt_file= params.pt_file, file_name= params.pt_file_name_net2, mess_out= True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_fea(params, xy, dl):\n",
    "    # mask\n",
    "    xys= []\n",
    "    for i in range(xy.shape[0]):\n",
    "        xys.append(torch.tensor([[i, 0, 0, 1209+ xy[i, 1]], [i, 0, 1, xy[i, 0]], [i, 1, 0, 1209+ xy[i, 1]], [i, 1, 1, xy[i, 0]], [i, 2, 0, 1209+ xy[i, 1]], [i, 2, 1, xy[i, 0]]]))\n",
    "    xys= torch.cat(xys, dim= 0)\n",
    "    # construct feature\n",
    "    a= torch.cat([dl.drug_struct_simi_mat[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.drug_micr_asso_mat_zy[:, xy[:, 1]].T.unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2)\n",
    "    # this feature has bad effect, so we mask it.\n",
    "    c= torch.cat([0* dl.drug_asso_simi_mat[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.drug_micr_asso_mat_zy[:, xy[:, 1]].T.unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2)\n",
    "    b= torch.cat([dl.drug_inter_mat[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.drug_micr_asso_mat_zy[:, xy[:, 1]].T.unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2)\n",
    "    d= torch.cat([dl.drug_micr_asso_mat_zy[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.micr_inter_mat[xy[:, 1], :].unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2) \n",
    "    e= torch.cat([dl.drug_micr_asso_mat_zy[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.micr_asso_simi_mat[xy[:, 1], :].unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2)\n",
    "    f= torch.cat([dl.drug_micr_asso_mat_zy[xy[:, 0], :].unsqueeze(dim= 1).unsqueeze(dim= 1), dl.micr_inte_simi_mat[xy[:, 1], :].unsqueeze(dim= 1).unsqueeze(dim= 1)], dim= 2)\n",
    "    fea= torch.cat([torch.cat([a, f], dim= -1), torch.cat([b, d], dim= -1), torch.cat([c, e], dim= -1)], dim= 1).to(torch.float)\n",
    "    fea[xys[:, 0], xys[:, 1], xys[:, 2], xys[:, 3]]= 0\n",
    "    return fea.to(params.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def avg_auc_aupr_cpt(test_xy, test_label, pred, ass_mat_shape):\n",
    "    label_mat, pred_mat= torch.zeros((ass_mat_shape)) -1, torch.zeros((ass_mat_shape)) -1\n",
    "    label_mat[test_xy[:, 0], test_xy[:, 1]], pred_mat[test_xy[:, 0], test_xy[:, 1]]= test_label* 1.0, pred\n",
    "    bool_mat4mark_test_examp= (label_mat!= -1)\n",
    "    aucs, auprs= [], []\n",
    "    for i in range(ass_mat_shape[0]):\n",
    "        test_examp_loc= bool_mat4mark_test_examp[i]\n",
    "        pos_num= label_mat[i, test_examp_loc].sum()\n",
    "        if pos_num> 0 and (test_examp_loc).sum()- pos_num> 0:\n",
    "            fpr4rowi, tpr4rowi, _= roc_curve(label_mat[i, test_examp_loc], pred_mat[i, test_examp_loc])\n",
    "            prec4rowi, recall4rowi, _= precision_recall_curve(label_mat[i, test_examp_loc], pred_mat[i, test_examp_loc])\n",
    "            prec4rowi[-1]= [1, 0][(int)(prec4rowi[-2]== 0)]\n",
    "            aucs.append(auc(fpr4rowi, tpr4rowi));auprs.append(auc(recall4rowi, prec4rowi))\n",
    "    return np.mean(aucs), np.mean(auprs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ep in range(params.epoch_net2):\n",
    "    for step, (xy, label) in enumerate(train_loader):\n",
    "        net.train()\n",
    "        xy, label= xy, label.to(params.device)\n",
    "        logp= net(xy[:, 0], xy[:, 1], get_fea(params, xy, dl))\n",
    "        loss= net.loss(logp, label)\n",
    "        optimizer.zero_grad();loss.backward();optimizer.step()\n",
    "        if step% 500== 0:print(f'epoch: {ep+ 1}, step: {step+ 1}, train loss: {loss}, acc: {(torch.max(logp, dim= 1)[1]== label).float().sum()* 1.0/ len(xy)}')\n",
    "    net.eval(); val_loss= 0; pred= []\n",
    "    with torch.no_grad():\n",
    "        for step, (xy, label) in enumerate(valid_loader):\n",
    "            xy, label= xy, label.to(params.device)\n",
    "            logp= net(xy[:, 0], xy[:, 1], get_fea(params, xy, dl))\n",
    "            val_loss= net.loss(logp, label).item()\n",
    "            pred.append(logp)\n",
    "    pred= torch.cat(pred, dim= 0).cpu()\n",
    "    roc_auc, aupr_auc= avg_auc_aupr_cpt(dl.valid_xy, valid_label, pred[:, 1], (1209, 172))\n",
    "    print(f'epoch: {ep+ 1}, valid loss: {val_loss}, acc: {(torch.max(pred, dim= 1)[1]== valid_label).float().sum()* 1.0/ len(valid_label)}, auc: {roc_auc}, aupr: {aupr_auc}')\n",
    "    earlystopping(-(roc_auc+ aupr_auc), net)\n",
    "    if earlystopping.flag== True:print(f'early_stopping');break;\n",
    "net= torch.load(f'{params.pt_file}//{params.pt_file_name_net2}')\n",
    "net.eval(); pred= []\n",
    "with torch.no_grad():\n",
    "    for step, (xy, label) in enumerate(test_loader):\n",
    "        xy, label= xy, label.to(params.device)\n",
    "        logp= net(xy[:, 0], xy[:, 1], get_fea(params, xy, dl))\n",
    "        pred.append(logp)\n",
    "    pred= torch.cat(pred, dim= 0).cpu()\n",
    "    precision, recall, threshold= precision_recall_curve(test_label, pred[:, 1])\n",
    "    roc_auc, aupr= roc_auc_score(test_label, pred[:, 1]), auc(recall, precision)\n",
    "    print(avg_auc_aupr_cpt(dl.test_xy, test_label, pred[:, 1], (1209, 172)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gj004",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
